import shutil
from GetFileLocation import *
from numpy import pi, cos, sin, arcsin
from obspy import UTCDateTime
from obspy.clients.fdsn import Client
from obspy.clients.fdsn.mass_downloader import Domain, Restrictions, MassDownloader


class Region(Domain):
    def __init__(self, event_lat, event_lon, minradius, maxradius, sllat, sllon, sulat, sulon):
        Domain.__init__(self)
        self.sllat = sllat
        self.sllon = sllon
        self.sulat = sulat
        self.sulon = sulon
        self.event_lat = event_lat
        self.event_lon = event_lon
        self.minradius = minradius
        self.maxradius = maxradius

    def conlat(self,latitude,longitude):
        lat = self.event_lat * pi / 180
        lon = self.event_lon * pi / 180
        lat0 = latitude * pi / 180
        lon0 = longitude * pi / 180
        s = sin(lat) * sin(lat0) + cos(lat) * cos(lat0) * cos(lon - lon0)
        clat = arcsin(s)
        clat = (pi / 2 - clat) * 180 / pi
        return clat

    # lon:-24~-12 lat:63~67
    def get_query_parameters(self):
        return {"minlatitude": self.sllat,
                "maxlatitude": self.sulat,
                "minlongitude": self.sllon,
                "maxlongitude": self.sulon}

    def is_in_domain(self, latitude, longitude):
        clat=self.conlat(latitude,longitude)
        if all((clat >= self.minradius, clat < self.maxradius)):
            return True
        return False


class RetrieveWaveform(object):
    def __init__(self, startTime, endTime, mmin, mmax,
                 sllat, sllon, sulat, sulon,
                 minradius, maxradius, projectLocation="DATA"):
        self.projectLocation = projectLocation
        self.starttime = startTime
        self.endtime = endTime
        self.mmin = mmin
        self.mmax = mmax
        self.sllat = sllat
        self.sllon = sllon
        self.sulat = sulat
        self.sulon = sulon
        self.minradius = minradius
        self.maxradius = maxradius
        self.ievent = 1
        self.isdownloading=False

    def _getEventsList(self):
        client = Client("IRIS")
        events = client.get_events(minmagnitude=self.mmin, maxmagnitude=self.mmax,
                                   starttime=self.starttime,
                                   endtime=self.endtime)
        # del self.starttime,self.endtime,self.mmax,self.mmin
        return events

    events = property(fget=_getEventsList)

    def _getEventsNum(self):
        return len(self.events)
    nevents = property(fget=_getEventsNum)

    def _getCatalogStorageLocation(self):
        if not os.path.exists(self.projectLocation):
            os.mkdir(self.projectLocation)
        return self.projectLocation + '/CATALOG_M%3.1f_M%3.1f_' % (self.mmin, self.mmax) + str(
            self.starttime.year) + '.TXT'


    catalogFile = property(fget=_getCatalogStorageLocation)

    def downloadWavefrom(self, event):
        if all([self.ievent == 1, os.path.exists(self.projectLocation),not self.isdownloading]):
            print("该项目已存在，新下载的数据将覆盖原项目")
            response=input("是否覆盖？1表示覆盖，输入其他内容表示不覆盖:\n")
            if response=="1":
                shutil.rmtree(self.projectLocation)
            else:
                return False
        self.isdownloading=True
        ctlg = open(self.catalogFile, 'a+')
        starttime = event.origins[0].time
        endtime = starttime + 1400.0
        latitude = event.origins[0].latitude
        longitude = event.origins[0].longitude
        depth = event.origins[0].depth
        if depth is None:
            return None
        if (depth > 50000.0):
            return None
        domain = Region(event_lat=latitude, event_lon=longitude
                        , minradius=30.00, maxradius=95.00
                        , sllat=self.sllat, sllon=self.sllon,
                        sulat=self.sulat, sulon=self.sulon)
        restrictions = Restrictions(
            starttime=starttime,
            endtime=endtime,
            reject_channels_with_gaps=True,
            channel="BH*",
            minimum_length=0.95,
            minimum_interstation_distance_in_m=0,
        )
        mdl = MassDownloader(providers=["IRIS"])
        event_dir = "%04d%02d%02d%03dT%02d%02d%02d" % (
            starttime.year, starttime.month, starttime.day, starttime.julday, starttime.hour, starttime.minute,
            starttime.second)
        mseed_storage = self.projectLocation + '/SEISMOGRAM/' + event_dir + '/MSEED'
        stationxml_storage = self.projectLocation + '/SEISMOGRAM/' + event_dir + '/XML'

        mdl.download(domain, restrictions,
                     mseed_storage=mseed_storage,
                     stationxml_storage=stationxml_storage)
        if os.path.exists(mseed_storage):
            event_string = "%06d|%s|%04d,%03d,%s/%s/%s/%s|%s %s %s %s %s" % (self.ievent,
                                                                             event_dir,
                                                                             starttime.year, starttime.julday,
                                                                             str(starttime)[11:22],
                                                                             str(latitude),
                                                                             str(longitude),
                                                                             str(depth / 1000.0),
                                                                             str(latitude),
                                                                             str(longitude),
                                                                             str(depth / 1000.0),
                                                                             str(event.magnitudes[0].mag),
                                                                             str(event.magnitudes[
                                                                                     0].magnitude_type))
            ctlg.write(event_string + "\n")
            qxml_storage = self.projectLocation + '/SEISMOGRAM/' + event_dir
            event.write(qxml_storage + "/" + event_dir + ".xml", format="QUAKEML")
            self.ievent += 1
        ctlg.close()

    def downloadAll(self):
        for event in self.events:
            status=self.downloadWavefrom(event)
            if status==False:
                break
        file=GetFileLocation(self.projectLocation)
        file.fileCheck()
        self.isdownloading = False


def downloadData(sllat = 63., sulat = 67., sllon = -24., sulon = -12., minradius = 30., maxradius = 90.):
    name = input("ProjectName:    ")
    year = input("Year:           ")
    mmin = float(input('Min Magnitude:  '))
    mmax = float(input('Max Magnitude:  '))
    startdate = UTCDateTime(year + '-01-01T00:00:00')
    enddate = UTCDateTime(year + '-12-31T23:59:59')
    download = RetrieveWaveform(startdate, enddate, mmin, mmax
                                , sllat, sllon, sulat, sulon
                                , minradius, maxradius
                                , projectLocation=name)
    download.downloadAll()
